package com.googlecode.jstructs.instr;

import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Method;
import java.security.ProtectionDomain;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.regex.Pattern;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.commons.EmptyVisitor;
import org.objectweb.asm.util.CheckClassAdapter;

public class StructInstrumentation implements ClassFileTransformer {

	ConcurrentHashMap<String, byte[]> newClasses = new ConcurrentHashMap<String, byte[]>();
	// TODO: Make weak...
	ConcurrentHashMap<String, Class<?>> definedClasses = new ConcurrentHashMap<String, Class<?>>();

	ConcurrentHashMap<String, Boolean> isValueType = new ConcurrentHashMap<String, Boolean>();

	private static Instrumentation instrumentation;

	private Method declaredMethod;

	@Override
	public byte[] transform(ClassLoader loader, String className,
			Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
			byte[] classfileBuffer) throws IllegalClassFormatException {
		TransformResult tr = transformClass(loader, classfileBuffer);
		if (tr == null) {
			return null;
		}
		for (Map.Entry<String, byte[]> e : tr.createdClasses.entrySet()) {
			loadClass(loader, e.getKey().replace('/', '.'), e.getValue());
		}
		return tr.transformedClassFileBuffer;
	}

	public TransformResult transformClass(ClassLoader loader,
			byte[] classfileBuffer) {
		try {
			ClassReader cr = new ClassReader(classfileBuffer);

			ClassTaster taster = new ClassTaster();
			taster.taste(cr, loader, this);
			isValueType.put(taster.name, taster.isValueType);
			if (taster.errSuperIsValueType) {
				System.err.println("Class " + taster.name
						+ " should not subclass " + cr.getSuperName());
			}
			if (!taster.isShouldInstrument()) {
				return null;
			}
			System.out.println(taster.name);
			TransformResult tr = new TransformResult();
			tr.className = taster.name;

			ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS
					| ClassWriter.COMPUTE_FRAMES);
			JStructClassInstr cv = new JStructClassInstr(this, loader, cw);
			cr.accept(cv, 0);
			tr.transformedClassFileBuffer = cw.toByteArray();
			checkClass(tr.transformedClassFileBuffer);
			// ASMifierClassVisitor asmifierClassVisitor = new
			// ASMifierClassVisitor(
			// new PrintWriter(System.out));
			// new ClassReader(tr.transformedClassFileBuffer).accept(
			// asmifierClassVisitor, 0);
			tr.createdClasses = cv.createdClasses;
			return tr;
		} catch (Exception ex) {
			ex.printStackTrace(System.err);
			throw new RuntimeException(ex);
		}
	}

	private void checkClass(byte[] byteArray) {
		try {
			ClassReader cr2 = new ClassReader(byteArray);
			ClassVisitor cv2 = new CheckClassAdapter(new EmptyVisitor());
			cr2.accept(cv2, 0);
		} catch (Exception ex) {
			ex.printStackTrace(System.err);
		}
	}

	static Pattern pattern = Pattern.compile("L?\\[*(.*);?");

	boolean isValueType(ClassLoader loader, String className) {
		Boolean b = isValueType.get(className);
		if (b != null) {
			return b.booleanValue();
		}

		String otherName = null;
		if ((className.startsWith("[") || className.startsWith("L"))
				&& className.endsWith(";")) {
			otherName = className;
			className = pattern.matcher(className).replaceAll("$1");
		}
		if (loader == null) {
			loader = Thread.currentThread().getContextClassLoader();
			if (loader == null) {
				loader = getClass().getClassLoader();
			}
		}
		InputStream in = loader.getResourceAsStream(className + ".class");
		if (in == null) {
			return false;
		}
		try {
			ClassReader r = new ClassReader(in);
			b = ClassTaster.VALUETYPE_CNAME.equals(r.getSuperName());
		} catch (IOException ex) {
			System.err.println("Ignoring: " + ex);
			b = Boolean.FALSE;
		}
		isValueType.put(className, b);
		if (otherName != null)
			isValueType.put(otherName, b);
		return b.booleanValue();
	}

	private void loadClass(ClassLoader loader, String className,
			byte[] byteArray) {
		checkClass(byteArray);
		if (instrumentation != null) {
			try {
				if (!definedClasses.containsKey(className)) {
					if (declaredMethod == null) {
						Method dm = ClassLoader.class.getDeclaredMethod(
								"defineClass", String.class, byte[].class,
								int.class, int.class);
						// ehehe
						dm.setAccessible(true);
						declaredMethod = dm;
					}
					Class<?> clazz = (Class<?>) declaredMethod.invoke(loader,
							className, byteArray, 0, byteArray.length);
					definedClasses.put(className, clazz);
				} else {
					instrumentation.redefineClasses(new ClassDefinition(
							definedClasses.get(className), byteArray));
				}
			} catch (Exception e) {
				throw new RuntimeException(
						"Error instrumenting continuation while loading generated class: "
								+ className, e);
			}
		} else {
			newClasses.put(className, byteArray);
		}
	}

	public ConcurrentHashMap<String, byte[]> getNewClasses() {
		return newClasses;
	}

	public static void premain(String agentArguments,
			Instrumentation instrumentation) {
		// System.out.println("jstructs instrumentation!");
		StructInstrumentation.instrumentation = instrumentation;
		instrumentation.addTransformer(new StructInstrumentation());
	}

}
